/*
** Copyright (C) 2010 Calin Radoni
**
** This file is part of CHScanner.
**
** CHScanner is free software: you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation, either version 3 of the License, or
** (at your option) any later version.
**
** CHScanner is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with CHScanner. If not, see <http://www.gnu.org/licenses/>.
*/

#include "StdAfx.h"
#include ".\scanmethodarp.h"

//--------------------------------------------------------------------

ScanMethodArp::ScanMethodArp(void) : ScanMethod()
{
	myType = __wmtArp;

	gwSearch = false;
	memset(&gwMAC[0], 0, 6);

	dnsSearch = false;
	memset(&dnsMAC[0], 0, 6);

	gwIP = 0; // will be completed in the "InternalInit" function
	dnsIP = 0; // will be completed in the "InternalInit" function

	localHostIP = 0x0100007F; // 127.0.0.1

	localIPaddr = 0; // will be completed in the "InternalInit" function
	localIPbcastaddr = 0; // will be completed in the "InternalInit" function
}

ScanMethodArp::~ScanMethodArp(void)
{
}

//--------------------------------------------------------------------

int ScanMethodArp::FindMacForGw(void)
{
	ScanPacket* pack;
	MArpCacheRec arpCacheRec;

	if(gwIP==0){
		// no gateway in this network !
		ClearMacAddress(gwMAC);
		return MSR_ERR_GW_MAC_NOT_FOUND;
	}

	if(config->xArpCacheOp!=__acoIgnore && config->xArpCacheOp!=__acoUpdOnly){
		if(arpCacheList!=NULL){
			if(arpCacheList->GetRecordByIPv4(winPcap->GetGatewayIP(), &arpCacheRec)){
				if(!arpCacheRec.addrMAC.IsEmpty()){
					gwMAC[0] = arpCacheRec.addrMAC.addr[0];
					gwMAC[1] = arpCacheRec.addrMAC.addr[1];
					gwMAC[2] = arpCacheRec.addrMAC.addr[2];
					gwMAC[3] = arpCacheRec.addrMAC.addr[3];
					gwMAC[4] = arpCacheRec.addrMAC.addr[4];
					gwMAC[5] = arpCacheRec.addrMAC.addr[5];
					return MSR_ERR_OK;
				}
			}
		}
	}

	// wait until everything is clear
	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}

	pack = FindFirstUnusedPacket();
	while(pack==NULL){
		SleepEx(1, TRUE);
		CheckResponses();
		CheckDeaths();
		pack = FindFirstUnusedPacket();
	}

	if(pack==NULL) return MSR_ERR_FAILED;
		
	pack->Initialize(true);
	pack->sTime       = GetTickCount();
	pack->addr        = gwIP;
	pack->port        = 0;
	pack->methodType  = __wmtArp;

	if(!SendArpRequest(gwIP)){
		ClearMacAddress(gwMAC);
		return MSR_ERR_GW_MAC_NOT_FOUND;
	}

	gwSearch = true;
	// wait until the end
	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}
	gwSearch = false;

	if(gwMAC[0]!=0x00) return MSR_ERR_OK;
	if(gwMAC[1]!=0x00) return MSR_ERR_OK;
	if(gwMAC[2]!=0x00) return MSR_ERR_OK;
	if(gwMAC[3]!=0x00) return MSR_ERR_OK;
	if(gwMAC[4]!=0x00) return MSR_ERR_OK;
	if(gwMAC[5]!=0x00) return MSR_ERR_OK;

	return MSR_ERR_GW_MAC_NOT_FOUND;
}

//--------------------------------------------------------------------

int ScanMethodArp::FindMacForDns(void)
{
	ScanPacket* pack;
	MArpCacheRec arpCacheRec;

	if(dnsIP==0){
		// no DNS in this network !
		ClearMacAddress(dnsMAC);
		return MSR_ERR_DNS_MAC_NOT_FOUND;
	}

	if(dnsIP==gwIP){
		for(int i=0; i<6; i++)
			dnsMAC[i] = gwMAC[i];
		return MSR_ERR_OK;
	}

	if(!winPcap->IpAddrBelongsToLocalNet(dnsIP)){
		memcpy(&dnsMAC[0], &gwMAC[0], 6);
		if(dnsMAC[0]!=0x00) return MSR_ERR_OK;
		if(dnsMAC[1]!=0x00) return MSR_ERR_OK;
		if(dnsMAC[2]!=0x00) return MSR_ERR_OK;
		if(dnsMAC[3]!=0x00) return MSR_ERR_OK;
		if(dnsMAC[4]!=0x00) return MSR_ERR_OK;
		if(dnsMAC[5]!=0x00) return MSR_ERR_OK;
		return MSR_ERR_DNS_MAC_NOT_FOUND;
	}

	if(config->xArpCacheOp!=__acoIgnore && config->xArpCacheOp!=__acoUpdOnly){
		if(arpCacheList!=NULL){
			if(arpCacheList->GetRecordByIPv4(dnsIP, &arpCacheRec)){
				if(!arpCacheRec.addrMAC.IsEmpty()){
					dnsMAC[0] = arpCacheRec.addrMAC.addr[0];
					dnsMAC[1] = arpCacheRec.addrMAC.addr[1];
					dnsMAC[2] = arpCacheRec.addrMAC.addr[2];
					dnsMAC[3] = arpCacheRec.addrMAC.addr[3];
					dnsMAC[4] = arpCacheRec.addrMAC.addr[4];
					dnsMAC[5] = arpCacheRec.addrMAC.addr[5];
					return MSR_ERR_OK;
				}
			}
		}
	}

	// wait until everything is clear
	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}

	pack = FindFirstUnusedPacket();
	while(pack==NULL){
		SleepEx(1, TRUE);
		CheckResponses();
		CheckDeaths();
		pack = FindFirstUnusedPacket();
	}

	if(pack==NULL) return MSR_ERR_FAILED;
		
	pack->Initialize(true);
	pack->sTime       = GetTickCount();
	pack->addr        = dnsIP;
	pack->port        = 0;
	pack->methodType  = __wmtArp;

	if(!SendArpRequest(dnsIP)){
		ClearMacAddress(dnsMAC);
		return MSR_ERR_DNS_MAC_NOT_FOUND;
	}

	dnsSearch = true;
	// wait until the end
	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}
	dnsSearch = false;

	if(dnsMAC[0]!=0x00) return MSR_ERR_OK;
	if(dnsMAC[1]!=0x00) return MSR_ERR_OK;
	if(dnsMAC[2]!=0x00) return MSR_ERR_OK;
	if(dnsMAC[3]!=0x00) return MSR_ERR_OK;
	if(dnsMAC[4]!=0x00) return MSR_ERR_OK;
	if(dnsMAC[5]!=0x00) return MSR_ERR_OK;

	return MSR_ERR_DNS_MAC_NOT_FOUND;
}

//--------------------------------------------------------------------

int ScanMethodArp::InternalInit(void)
{
	gwSearch = false;
	memset(&gwMAC[0], 0, 6);

	dnsSearch = false;
	memset(&dnsMAC[0], 0, 6);

	if(workPhase!=__wptArp) return MSR_ERR_PHASE;

	gwIP  = winPcap->GetGatewayIP();
	dnsIP = winPcap->GetDnsServerIP();

	localIPaddr = winPcap->GetLocalIP();
	localIPbcastaddr = winPcap->GetLocalIPBroadcastAddress();

	return MSR_ERR_OK;
}

//--------------------------------------------------------------------

bool ScanMethodArp::FillMacFromCache(MFullAddress* addrIn)
{
	DWORD theIpAddr;

	if(addrIn==NULL) return false;

	theIpAddr = addrIn->addrIP.GetD();

	// localhost address ?
	if(theIpAddr==localHostIP){
		winPcap->FillSrcMAC(addrIn->addrMAC.addr);
		addrIn->gwMAC = false;
		addrIn->alive = false;
		return true;
	}

	// local IP address ?
	if(theIpAddr==localIPaddr){
		winPcap->FillSrcMAC(addrIn->addrMAC.addr);
		addrIn->gwMAC = false;
		addrIn->alive = true;
		return true;
	}

	// local network broadcast address ?
	if(theIpAddr==localIPbcastaddr){
		addrIn->addrMAC.SetAsBroadcastAddress();
		addrIn->gwMAC = false;
		addrIn->alive = true;
		return true;
	}

	// gateway's address ?
	if(theIpAddr==gwIP){
		addrIn->addrMAC.LoadFromBytePtr(&gwMAC[0]);
		addrIn->gwMAC = false;
		addrIn->alive = true;
		return true;
	}

	// address of DNS ?
	if(theIpAddr==dnsIP){
		if(winPcap->IpAddrBelongsToLocalNet(dnsIP)){
			addrIn->addrMAC.LoadFromBytePtr(&dnsMAC[0]);
			addrIn->gwMAC = false;
			addrIn->alive = true;
		}
		else{
			addrIn->addrMAC.LoadFromBytePtr(&dnsMAC[0]);
			addrIn->gwMAC = true;
			addrIn->alive = false;
		}
		return true;
	}

	MArpCacheRec arpCacheRec;

	if(config->xArpCacheOp!=__acoIgnore && config->xArpCacheOp!=__acoUpdOnly){
		if(arpCacheList!=NULL){
			if(arpCacheList->GetRecordByIPv4(theIpAddr, &arpCacheRec)){
				if(!arpCacheRec.addrMAC.IsEmpty()){
					addrIn->addrMAC.LoadFromBytePtr(&arpCacheRec.addrMAC.addr[0]);
					addrIn->gwMAC = false;
					addrIn->alive = false;
					return true;
				}
			}
		}
	}

	if(!winPcap->IpAddrBelongsToLocalNet(theIpAddr)){
		addrIn->addrMAC.LoadFromBytePtr(&gwMAC[0]);
		addrIn->gwMAC = true;
		addrIn->alive = false;
		return true;
	}

	return false;
}

//--------------------------------------------------------------------

int ScanMethodArp::InternalExecute(void)
{
	int i, cnt, res;
	ScanPacket* pack;
	MFullAddress* addr;

	if(addrList==NULL) return MSR_ERR_PARAMETER;

	res = FindMacForGw();
	if(res==MSR_ERR_GW_MAC_NOT_FOUND)
		res = MSR_ERR_OK;
	if(res!=MSR_ERR_OK) return res;

	res = FindMacForDns();
	if(res==MSR_ERR_DNS_MAC_NOT_FOUND)
		res = MSR_ERR_OK;
	if(res!=MSR_ERR_OK) return res;

	cnt = addrList->GetSize();
	for(i=0; i<cnt && !stopStopStop; i++){
		if(PacketsUsed()>0){
			CheckResponses();
			CheckDeaths();
		}

		addr = addrList->GetByIdx(i);
		if(!addr->addrIP.IsEmpty()){
			// it is an IPv4 address
			if(!FillMacFromCache(addr)){
				pack = FindFirstUnusedPacket();
				while(pack==NULL){
					SleepEx(waitTime, TRUE);
					CheckResponses();
					CheckDeaths();
					pack = FindFirstUnusedPacket();
				}

				if(pack!=NULL && addr!=NULL){
					pack->Initialize(true);
					pack->sTime       = GetTickCount();
					pack->addr        = addr->addrIP.GetD();
					pack->port        = 0;
					pack->methodType  = myType;
					pack->fullAddress = addr;

					if(!SendArpRequest(pack->addr))
						pack->Initialize(false);

					if((packSent % notifTcp) == 0 && packSent>0)
						if(hWnd!=NULL)
							PostMessage(hWnd, CHS_MSG_PACKETS_SENT, 0, (LPARAM)packSent);
				}
			}
		}
	}

	// wait until the end
	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}

/*
	CString str;
	int good, bad;
	good = 0;
	bad = 0;
	for(i=0; i<cnt; i++){
		addr = addrList->GetItem(i);
		if(!addr->addrMAC.IsEmpty()){

			in_addr ia;
			ia.S_un.S_addr = addr->addrIP;
			str.Format(_T("   > %d.%d.%d.%d = %02x-%02x-%02x-%02x-%02x-%02x\n"),
				ia.S_un.S_un_b.s_b1, ia.S_un.S_un_b.s_b2, ia.S_un.S_un_b.s_b3, ia.S_un.S_un_b.s_b4,
				addr->addrMAC[0], addr->addrMAC[1], addr->addrMAC[2], addr->addrMAC[3], addr->addrMAC[4], addr->addrMAC[5]);
			OutputDebugString(str);

			good++;
		}
		else{
			bad++;
		}
	}
	str.Format(_T("Found %d from a total of %d\n"), good, good + bad);
	OutputDebugString(str);
*/
	return MSR_ERR_OK;
}
//---------------------------------------------------------------------------

bool ScanMethodArp::SendArpRequest(DWORD dstIpIn)
{
	BYTE srcMAC[6];
	BYTE srcIP[4];
	BYTE dstIP[4];
	in_addr inAddr;

	winPcap->FillSrcMAC(&srcMAC[0]);

	winPcap->FillSrcIp(&srcIP[0]);

	inAddr.S_un.S_addr = dstIpIn;
	dstIP[0] = inAddr.S_un.S_un_b.s_b1;
	dstIP[1] = inAddr.S_un.S_un_b.s_b2;
	dstIP[2] = inAddr.S_un.S_un_b.s_b3;
	dstIP[3] = inAddr.S_un.S_un_b.s_b4;

	if(!packetBuilderEx->BuildPacketArpRequest(srcMAC, srcIP, dstIP))
		return false;

	return SendPacket(packetBuilderEx->GetPacket(), (int)packetBuilderEx->GetPacketSize());
}

//---------------------------------------------------------------------------

void ScanMethodArp::ProcessArpReply(BYTE* dataIn, int sizeIn)
{
	in_addr inAddr;
	ScanPacket* pack;
	MFullAddress* addr;

	if(sizeIn<42){
		// nasol, pachet prea mic
		return;
	}

	inAddr.S_un.S_un_b.s_b1 = dataIn[14];
	inAddr.S_un.S_un_b.s_b2 = dataIn[15];
	inAddr.S_un.S_un_b.s_b3 = dataIn[16];
	inAddr.S_un.S_un_b.s_b4 = dataIn[17];

	pack = FindPacket(inAddr.S_un.S_addr, 0, myType);
	if(pack==NULL) return;

	if(gwSearch){
		if(inAddr.S_un.S_addr==winPcap->GetGatewayIP()){
			memcpy(&gwMAC[0], &dataIn[8], 6);
			pack->Initialize(false);
		}
		return;
	}

	if(dnsSearch){
		if(inAddr.S_un.S_addr==winPcap->GetDnsServerIP()){
			memcpy(&dnsMAC[0], &dataIn[8], 6);
			pack->Initialize(false);
		}
		return;
	}

	addr = pack->fullAddress;
	if(addr!=NULL){
		addr->addrMAC.LoadFromBytePtr(&dataIn[8]);
		addr->gwMAC = false;
		addr->alive = true;
	}
	pack->Initialize(false);

//	CString str;
//	str.Format(_T("   > Received ARP Replay for %d.%d.%d.%d: %02x-%02x-%02x-%02x-%02x-%02x\n"),
//					dataIn[14], dataIn[15], dataIn[16], dataIn[17],
//					dataIn[8], dataIn[9], dataIn[10], dataIn[11], dataIn[12], dataIn[13]);
//	OutputDebugString(str);
}

//---------------------------------------------------------------------------

void ScanMethodArp::FillDnsServerMAC(BYTE* dataIn)
{
	if(dataIn==NULL) return;
	memcpy(dataIn, &dnsMAC[0], 6);
}

//---------------------------------------------------------------------------

void ScanMethodArp::ClearMacAddress(BYTE* dataIn)
{
	if(dataIn==NULL) return;

	for(int i=0; i<6; i++)
		dataIn[i] = 0;
}

//---------------------------------------------------------------------------
